import sys
import torch
import torch.nn as nn
import torch.nn.functional as F
from math import sqrt
import numpy as np
from orbitP.script import config


class FFN(nn.Module):
    def __init__(self, in_dim, hidden_dim, dropout):
        super().__init__()
        self.norm1 = nn.LayerNorm(in_dim)
        self.norm2 = nn.LayerNorm(in_dim)
        self.fc1 = nn.Linear(in_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, in_dim)
        self.activation = nn.GELU()
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        y = self.norm1(x)
        y = self.fc2(self.activation(self.fc1(y)))
        y = self.dropout(y)
        return self.norm2(x + y)  # 残差连接


class TCM(nn.Module):
    """沿着时间维度(L)的MLP，使用 ReplicationPad1d"""
    def __init__(self,in_dim, hidden_dim, kernel_size, dropout):
        super().__init__()
        self.norm1 = nn.LayerNorm(in_dim)
        self.norm2 = nn.LayerNorm(in_dim)
        self.pad1 = nn.ReplicationPad1d(((kernel_size-1)//2, (kernel_size-1)//2))
        self.pad2 = nn.ReplicationPad1d(((kernel_size-1)//2, (kernel_size-1)//2))
        self.conv1 = nn.Conv1d(in_channels=in_dim, out_channels=hidden_dim, kernel_size=kernel_size)
        self.conv2 = nn.Conv1d(in_channels=hidden_dim, out_channels=in_dim, kernel_size=kernel_size)
        self.activation = nn.GELU()
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        # 输入: [B, L, D] -> Conv1d 需要 [B, D, L]
        y = x = self.norm1(x)

        y = y.transpose(-1, 1)         # B,D,L
        y = self.pad1(y)               # B,D,L + pad
        y = self.conv1(y)              # B,2D,L
        y = self.activation(y)
        y = self.dropout(y)

        y = self.pad2(y)               # 再复制填充
        y = self.conv2(y)              # B,D,L
        y = self.dropout(y)
        y = y.transpose(-1, 1)         # B,L,D
        return self.norm2(x + y)       # 残差


class SingleEncoder(nn.Module):
    """单一特征组编码器：特征维度 + 时间维度"""

    def __init__(self, in_dim, hidden_dim, kernel_size, dropout):
        super().__init__()
        self.ffn = FFN(in_dim, hidden_dim, dropout)
        self.tcm = TCM(in_dim,hidden_dim, kernel_size, dropout)

    def forward(self, x):
        # x: [B, L, in_dim]
        x = self.ffn(x)  # [B, L, hidden_dim]
        x = self.tcm(x)  # [B, L, hidden_dim]
        return x


class FeatureEncoder(nn.Module):
    def __init__(self, args):
        super().__init__()
        # 观测段输入
        self.embed_dim = args.embed_dim
        self.satEnc = nn.Embedding(config.GPS_size + 1, args.embed_dim)
        self.val_embed = SingleEncoder(in_dim=config.axis + 1, hidden_dim=args.head_dim, kernel_size=args.head_kernel_size, dropout=args.dropout)  # 位置误差特征
        self.physical_embed = SingleEncoder(in_dim=args.feature_size - (config.axis + 1 + config.GPS_stampSize),
                                            hidden_dim=args.head_dim, kernel_size=args.head_kernel_size,
                                            dropout=args.dropout)  # 物理信息特征
        self.sat_embed = SingleEncoder(in_dim=self.embed_dim, hidden_dim=args.head_dim,
                                        kernel_size=args.head_kernel_size, dropout=args.dropout)  # 时间戳特征
        self.time_embed = SingleEncoder(in_dim=config.stampSize, hidden_dim=args.head_dim,
                                        kernel_size=args.head_kernel_size, dropout=args.dropout)  # 时间戳特征
        # 预测阶段输入
        self.norm1 = nn.LayerNorm(args.feature_size-1+self.embed_dim+config.stampSize)
        self.norm2 = nn.LayerNorm(args.feature_size-1+self.embed_dim+config.stampSize)
        self.proj = SingleEncoder(in_dim=args.feature_size-1+self.embed_dim+config.stampSize, hidden_dim=args.head_dim, kernel_size=args.head_kernel_size, dropout=args.dropout)

    def forward(self, x_enc, x_mark_enc, x_dec, x_mark_dec):
        """
        输入:
            x_enc:  [B, L, feature_dim]    历史输入
            x_mark_enc: [B, L, time_dim]   历史时间特征
            x_dec:  [B, L, feature_dim]    预测目标输入
            x_mark_dec: [B, L, time_dim]   预测时间特征
        输出:
            enc_input: 编码器输入 [B, L, d_model//2]
            prd_input: 预测器输入 [B, L, d_model//2]
        """
        # ----------- 历史输入处理 -----------
        # 分离基础特征和卫星ID
        x_enc_pre = x_enc[:, :, :-config.GPS_stampSize]
        sat_id = x_enc[:, :, -config.GPS_stampSize:].long().squeeze(-1)
        sat_id = self.satEnc(sat_id)
        x = torch.cat((x_enc_pre,x_mark_enc,sat_id),dim=-1)
        # 各特征组编码
        val_enc = self.val_embed(x[:, :, :config.axis + 1])
        # [margin,np.sin(beta_rad),np.cos(beta_rad), Kp, F10.7]
        physical_enc = self.physical_embed(x[:, :, config.axis + 1:config.axis + config.physicalSize + 1])
        time_enc = self.time_embed(x[:, :, config.axis + config.physicalSize + 1:config.axis + config.physicalSize + config.stampSize + 1])
        sat_enc = self.sat_embed(x[:, :, -self.embed_dim:])
        y = torch.cat((val_enc, physical_enc, time_enc, sat_enc), dim=-1)
        y = self.norm1(y)
        y = self.proj(y)
        return self.norm2(x+y)